from tensorflow import keras


class MOneHotDecisionModelV2:
    def __init__(self,
                 learning_rate,
                 activation,
                 kernel_initializer,
                 epochs,
                 name,
                 min_shape
                 ):
        self.learning_rate = learning_rate
        self.activation = activation
        self.kernel_initializer = kernel_initializer
        self.epochs = epochs
        self.model = None
        self.name = name
        self.min_shape = min_shape

    def create_model(self):
        model = keras.Sequential(name=self.name)
        model.add(keras.layers.Input(shape=self.min_shape))
        model.add(keras.layers.LSTM(units=128, activation=keras.activations.sigmoid,
                                    kernel_initializer=keras.initializers.glorot_uniform,
                                    dropout=0.02,
                                    return_sequences=True))
        model.add(keras.layers.LayerNormalization())
        model.add(keras.layers.LSTM(units=64, activation=keras.activations.sigmoid,
                                    kernel_initializer=keras.initializers.glorot_normal))
        model.add(keras.layers.Dense(units=256,
                                     activation=keras.activations.sigmoid,
                                     kernel_initializer=keras.initializers.glorot_normal))
        model.add(keras.layers.Dropout(0.02))
        model.add(keras.layers.Dense(units=128,
                                     activation=keras.activations.sigmoid,
                                     kernel_initializer=keras.initializers.glorot_normal))
        model.add(keras.layers.LayerNormalization())
        model.add(keras.layers.Dense(units=1,
                                     name="output",
                                     activation=keras.activations.sigmoid))

        model.compile(optimizer=keras.optimizers.Adam(learning_rate=self.learning_rate, weight_decay=0.01),
                      loss=keras.losses.binary_crossentropy, metrics=["accuracy"])

        self.model = model
        return model

    def __generate_mode_path(self):
        return f'./model/{self.name}.kares'

    def load_model(self):
        print("load model" + self.__generate_mode_path())
        self.model = keras.models.load_model(self.__generate_mode_path())
        return self.model

    def save_model(self):
        self.model.save(self.__generate_mode_path())

    def train_model(self,
                    train_data,
                    train_target,
                    val_data,
                    val_target):
        path_checkpoint = self.__generate_mode_path()
        model_checkpoint_callback = keras.callbacks.ModelCheckpoint(
            monitor="val_loss",
            filepath=path_checkpoint,
            verbose=1,
            save_best_only=True,
        )
        history = self.model.fit(
            x=train_data,
            y=train_target,
            batch_size=2048,
            epochs=self.epochs,
            verbose=2,
            callbacks=model_checkpoint_callback,
            validation_data=(val_data, val_target),
        )
        return history
